# -*-coding=utf-8-*-

"""
构建一个简单势能面
"""

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


class SimpleSurface():
    def __init__(self):
        pass

    def get_value(self, x, y):
        """
        计算函数
        """
        return np.sin(x) + np.sin(y)

    def show_surface(self):
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        x_array = np.arange(-8, 8, 0.5)
        y_array = np.arange(-8, 8, 0.5)
        x, y = np.meshgrid(x_array, y_array)
        z = [self.get_value(x_, y_) for x_, y_ in zip(x, y)]
        ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap=plt.get_cmap('rainbow'))
        # ax.contourf(x, y, z, zdir='z', offset=2)
        ax.set_zlim(-4, 4)
        plt.show()


if __name__ == "__main__":
    ss = SimpleSurface()
    ss.show_surface()
